/*!
  \file
  Implementation of the FFT for HHVolumeGPU (uses CUFFT)
*/

#include <cuComplex.h>
#include <cufft.h>

#include "fftcuda.hpp"

// =========================================

#define CUFFT_SAFE_CALL( call ) do {            \
    cufftResult fftErr = call;                                  \
    if( CUFFT_SUCCESS != fftErr ) {                                     \
      fprintf( stderr, "CUFFT Error in file '%s' on line %i: %#x\n",    \
               __FILE__, __LINE__, fftErr );                            \
      exit( EXIT_FAILURE );                                             \
    } } while( 0 );

// =========================================

void FFTimageCUDA::Forward( SciGPU::HHV::HHVolume< std::complex<float>,2>& dst,
                              const SciGPU::HHV::HHVolume<float,2>& src ) const {
  SciGPU::HHV::HHVolumeGPU< float,2>::convertor srcConvert( src );
  SciGPU::HHV::HHVolumeGPU< std::complex<float>,2>::convertor dstConvert( dst, false );
  
  this->ForwardTransform( dstConvert.get(), srcConvert.getConst() );
}


__global__ void R2Ckernel( cuFloatComplex* dst,
                           const float* src,
                           const size_t nVals ) {
  const size_t idx = threadIdx.x + ( blockIdx.x * blockDim.x );
  
  if( idx < nVals ) {
    dst[idx] = make_cuFloatComplex( src[idx], 0 );
  }
}


void FFTimageCUDA::ForwardTransform( SciGPU::HHV::HHVolumeGPU< std::complex<float>,2>& dst,
                                     const SciGPU::HHV::HHVolumeGPU<float,2>& src ) const {
  const SciGPU::HHV::DimensionObject<2> dims = src.dimensions();
 
  SciGPU::HHV::HHVolumeGPU< std::complex<float>,2> srcComplex;
  srcComplex.Allocate( dims );
  dst.Allocate( dims );
    
  // Upconvert to float
  
  const unsigned int kBlockSize = 256;
  dim3 grid, threads;
  threads.x = kBlockSize;
  threads.y = threads.z = 1;
  grid.x = static_cast<int>( ceil( static_cast<float>( src.hyxels() ) / threads.x ) );
  grid.y = grid.z = 1;
  
  R2Ckernel<<<grid,threads>>>( reinterpret_cast<cuFloatComplex*>(srcComplex.start()),
                               src.start(),
                               src.hyxels() );
  cudaError_t err = cudaThreadSynchronize();
  if( cudaSuccess != err ) {
    throw std::runtime_error( "R2Ckernel failed" );
  }

  // Do CUFFT things
  cufftHandle plan;

  CUFFT_SAFE_CALL( cufftPlan2d( &plan, dims[0], dims[1], CUFFT_C2C ) );

  CUFFT_SAFE_CALL( cufftExecC2C( plan,
                                 reinterpret_cast<cuFloatComplex*>(srcComplex.start()),
                                 reinterpret_cast<cuFloatComplex*>(dst.start()),
                                 CUFFT_FORWARD ) );

  CUFFT_SAFE_CALL( cufftDestroy(plan) );

}


// ==========================================


void FFTimageCUDA::Inverse( SciGPU::HHV::HHVolume<float,2>& dst,
                            const SciGPU::HHV::HHVolume< std::complex<float>,2>& src ) const {
  SciGPU::HHV::HHVolumeGPU< float,2>::convertor dstConvert( dst, false );
  SciGPU::HHV::HHVolumeGPU< std::complex<float>,2>::convertor srcConvert( src );
  
  this->InverseTransform( dstConvert.get(), srcConvert.getConst() );
}



__global__ void RescaleC2Rkernel( float* dst,
                                  const cuFloatComplex* src,
                                  const float a,
                                  const size_t nVals ) {
  const size_t idx = threadIdx.x + ( blockIdx.x * blockDim.x );
  
  if( idx < nVals ) {
    dst[idx] = a * cuCrealf( src[idx] );
  }
}


void FFTimageCUDA::InverseTransform( SciGPU::HHV::HHVolumeGPU<float,2>& dst,
                                     const SciGPU::HHV::HHVolumeGPU< std::complex<float>,2>& src ) const {
  const SciGPU::HHV::DimensionObject<2> dims = src.dimensions();
  SciGPU::HHV::HHVolumeGPU< std::complex<float>,2> dstComplex;
  dstComplex.Allocate( dims );
  dst.Allocate( dims );
  
  // Do the transform
  cufftHandle plan;
  
  CUFFT_SAFE_CALL( cufftPlan2d( &plan, dims[0], dims[1], CUFFT_C2C ) );
  
  CUFFT_SAFE_CALL( cufftExecC2C( plan,
                                 reinterpret_cast<cuFloatComplex*>(const_cast<std::complex<float>*>(src.start())),
                                 reinterpret_cast<cuFloatComplex*>(dstComplex.start()),
                                 CUFFT_INVERSE ) );
  
  CUFFT_SAFE_CALL( cufftDestroy( plan ) );
  
  // Rescale and convert
  const unsigned int kBlockSize = 256;
  dim3 grid, threads;
  threads.x = kBlockSize;
  threads.y = threads.z = 1;
  grid.x = static_cast<int>( ceil( static_cast<float>( dst.hyxels() ) / threads.x ) );
  grid.y = grid.z = 1;
  
  RescaleC2Rkernel<<<grid,threads>>>( dst.start(),
                                      reinterpret_cast<cuFloatComplex*>(dstComplex.start()),
                                      1.0f/dst.hyxels(),
                                      dst.hyxels() );
  cudaError_t err = cudaThreadSynchronize();
  if( cudaSuccess != err ) {
    throw std::runtime_error( "RescaleC2Rkernel failed" );
  }
  
}
